from socket import socket, AF_INET, SO_REUSEADDR, SOCK_STREAM, SOL_SOCKET, timeout
from NNserver.myThread import FC_Para
from Data_Pack import myPack, myRecv
from NNserver.myDset import CFP_Dset, DL_BSIZE
from torch.utils.data import DataLoader
import torch
import numpy as np
import random
import os


class myServer():
    def __init__(self, args):
        # 随机数
        if args.seed != 0:
            torch.manual_seed(args.seed)
            torch.cuda.manual_seed_all(args.seed)
            np.random.seed(args.seed)
            random.seed(args.seed)
            torch.use_deterministic_algorithms(True)
            torch.backends.cudnn.benchmark = False
            os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
            os.environ['PYTHONHASHSEED'] = str(args.seed)
        # 构造工作线程
        self.myPa = FC_Para()
        self.work_on = False
        # 构造连接套接字
        self.sc_s = socket(AF_INET, SOCK_STREAM)
        self.wsec_s = 120.0
        self.sc_c = None
        self.wsec_c = 60.0
        self.addr_c = None
        # 构造指令-函数映射关系
        self.cmd_dic = {
            'close': self.Work_close,
            'loaddset': self.Work_ldset,
            'samget': self.Work_samget,
            'slwfile': self.Work_slwfile,
            'nnopa1': self.Work_nnopa1,
            'nnstop1': self.Work_nnstop1,
            'nngs1': self.Work_nngs1,
            'nngd1': self.Work_nngd1,
            'nnopa2': self.Work_nnopa2,
            'nnstop2': self.Work_nnstop2,
            'nngd2': self.Work_nngd2
        }
        self.cmd_fun = None
        self.tcp_text = None
        # 绑定地址，连接客户端
        self.Con_wait(args)

    def Con_wait(self, args):
        self.sc_s.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        self.sc_s.settimeout(self.wsec_s)
        try:
            self.sc_s.bind((args.host, args.port))
        except Exception as e:
            print(e)
            return
        self.sc_s.listen(1)
        print('--正在监听：{}'.format(self.sc_s.getsockname()))
        try:
            self.sc_c, self.addr_c = self.sc_s.accept()
        except Exception as e:
            print(e)
        else:
            print('收到连接：{}'.format(self.addr_c))
            self.sc_c.settimeout(self.wsec_c)
            self.work_on = True

    def Work(self):
        print('--正在等待客户端指令')
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.cmd_fun = self.cmd_dic[self.tcp_text['cmd']]
        except timeout as e:
            print(e)
            return
        except Exception as e:
            print(type(e), e)
            self.sc_c.close()
            self.work_on = False
            return
        print('收到指令：{}'.format(self.tcp_text['cmd']))
        self.cmd_fun()

    def Work_close(self):
        self.sc_c.close()
        self.work_on = False

    def seed_worker(worker_id):
        worker_seed = torch.initial_seed() % 2**32
        np.random.seed(worker_seed)
        random.seed(worker_seed)

    def Work_ldset(self):
        self.myPa.dset_path = self.tcp_text['droot']
        print('数据集加载中')
        try:
            self.myPa.dset_train = CFP_Dset(self.myPa.dset_path, train=True)
            self.myPa.dset_test = CFP_Dset(self.myPa.dset_path, train=False)
        except Exception as e:
            print(e)
            self.myPa.dset_on = False
            self.tcp_text = {'dset_on': False}
            print('数据集加载失败！')
        else:
            self.myPa.dload_train = DataLoader(self.myPa.dset_train, batch_size=DL_BSIZE, shuffle=True,
                                                    drop_last=False, worker_init_fn=self.seed_worker)
            self.myPa.dload_test = DataLoader(self.myPa.dset_test, batch_size=DL_BSIZE, shuffle=False,
                                                    drop_last=False, worker_init_fn=self.seed_worker)
            self.myPa.dset_on = True
            self.tcp_text = {'dset_on': True,
                            'strain': str(len(self.myPa.dset_train)),
                            'stest': str(len(self.myPa.dset_test)),
                            'ltrain': len(self.myPa.dset_train),
                            'ltest': len(self.myPa.dset_test)}
            print('数据集加载成功！')
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_samget(self):
        self.myPa.sam_index = self.tcp_text['index']
        self.myPa.sam_count = self.tcp_text['count']
        print('正在获取样本图')
        if self.myPa.sam_index == 0:
            self.myPa.sam_tsr, self.myPa.sam_tar = self.myPa.dset_train[self.myPa.sam_count]
        elif self.myPa.sam_index == 1:
            self.myPa.sam_tsr, self.myPa.sam_tar = self.myPa.dset_test[self.myPa.sam_count]
        self.myPa.sam_img = self.myPa.sam_trans(self.myPa.sam_tsr)
        self.myPa.sam_tsr = self.myPa.sam_tsr.cuda()
        self.myPa.sam_tar = torch.tensor(self.myPa.sam_tar).cuda()
        self.myPa.sam_str = self.myPa.dset_train.classes[self.myPa.sam_tar]
        self.tcp_text = {'img': self.myPa.sam_img, 'str': self.myPa.sam_str}
        print('样本图获取完成！')
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_slwfile(self):
        self.myPa.nn_wsave = self.tcp_text['nn_wsave']
        if self.myPa.nn_wsave:
            print('正在获取权重文件')
            self.myPa.nn_wfile = {'s_net': self.myPa.nn_net.state_dict(),
                                    's_opt': self.myPa.nn_opt.state_dict()}
            print('权重文件获取完成！')
        else:
            print('正在载入权重文件')
            self.myPa.nn_wfile = self.tcp_text['nn_wfile']
            self.myPa.nn_net.load_state_dict(self.myPa.nn_wfile['s_net'])
            self.myPa.nn_opt.load_state_dict(self.myPa.nn_wfile['s_opt'])
            print('权重文件载入完成！')
        self.tcp_text = {'nn_wfok': True}
        if self.myPa.nn_wsave:
            self.tcp_text['nn_wfile'] = self.myPa.nn_wfile
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nnopa1(self):
        nno1_ok = not self.myPa.nn_workon
        if nno1_ok:
            print('神经网络开始运作（模式1）')
            self.myPa.nno1_dset = self.tcp_text['nno1_dset']
            self.myPa.nno1_back = self.tcp_text['nno1_back']
            self.myPa.nno1_epoch = self.tcp_text['nno1_epoch']
            self.myPa.nno1_acur = self.tcp_text['nno1_acur']
            self.myPa.nno1_loss = self.tcp_text['nno1_loss']
            self.myPa.nn_workon = True
            self.myPa.nn_ostate = 1
            self.myPa.Th_run(1)
        else:
            print('神经网络正在运作')
        self.tcp_text = {'nno1_ok': nno1_ok,
                        'nn_workon': self.myPa.nn_workon,
                        'nn_ostate': self.myPa.nn_ostate}
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nnstop1(self):
        self.tcp_text = {'nn_ostate': self.myPa.nn_ostate}
        if self.myPa.nn_workon and (self.myPa.nn_ostate == 1):
            print('神经网络正在中止（模式1）')
            self.tcp_text['nns1_ok'] = True
            self.tcp_text['nn_workon'] = False
            self.myPa.nn_sn1 = True
        else:
            self.tcp_text['nns1_ok'] = False
            self.tcp_text['nn_workon'] = self.myPa.nn_workon
            if self.myPa.nn_workon:
                print('神经网络未在该模式')
            else:
                print('神经网络未在运行')
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nngs1(self):
        self.tcp_text = {'nn_workon': self.myPa.nn_workon,
                        'nn_ostate': self.myPa.nn_ostate,
                        'nn_nowepoch': self.myPa.nn_nowepoch,
                        'nn_nowbatch': self.myPa.nn_nowbatch}
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nngd1(self):
        self.tcp_text = {'nn_a1_ok': self.myPa.nn_a1_ok,
                        'nn_l1_ok': self.myPa.nn_l1_ok,
                        'a1_depoch': self.myPa.nno1_epoch}
        if self.myPa.nn_a1_ok:
            self.tcp_text['a1_nums'] = self.myPa.a1_nums
            self.tcp_text['a1_cors'] = self.myPa.a1_cors
            self.tcp_text['a1_rates'] = self.myPa.a1_rates
        if self.myPa.nn_l1_ok:
            self.tcp_text['a1_losses'] = self.myPa.a1_losses
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nnopa2(self):
        nno2_ok = not self.myPa.nn_workon
        if nno2_ok:
            print('神经网络开始运作（模式2）')
            self.myPa.nn_workon = True
            self.myPa.nn_ostate = 2
            self.myPa.Th_run(2)
        else:
            print('神经网络正在运作')
        self.tcp_text = {'nno2_ok': nno2_ok,
                        'nn_workon': self.myPa.nn_workon,
                        'nn_ostate': self.myPa.nn_ostate}
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nnstop2(self):
        self.tcp_text = {'nn_ostate': self.myPa.nn_ostate}
        if self.myPa.nn_workon and (self.myPa.nn_ostate == 2):
            print('神经网络正在中止（模式2）')
            self.tcp_text['nns2_ok'] = True
            self.tcp_text['nn_workon'] = False
            self.myPa.nn_sn2 = True
        else:
            self.tcp_text['nns2_ok'] = False
            self.tcp_text['nn_workon'] = self.myPa.nn_workon
            if self.myPa.nn_workon:
                print('神经网络未在该模式')
            else:
                print('神经网络未在运行')
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nngd2(self):
        self.tcp_text = {'nn_a2_ok': self.myPa.nn_a2_ok,
                        'nn_l2_ok': self.myPa.nn_l2_ok,
                        'a2_depoch': self.myPa.nn_nowepoch + 1}
        if self.myPa.nn_a2_ok:
            self.tcp_text['a2_nums'] = self.myPa.a2_nums
            self.tcp_text['a2_cors'] = self.myPa.a2_cors
            self.tcp_text['a2_rates'] = self.myPa.a2_rates
        if self.myPa.nn_l2_ok:
            self.tcp_text['a2_losses'] = self.myPa.a2_losses
        self.sc_c.sendall(myPack(self.tcp_text))
